package svsim

import java.io.{BufferedReader, BufferedWriter, File, FileWriter, InputStreamReader, PrintWriter}
import java.nio.file.attribute.BasicFileAttributes
import java.nio.file.{FileVisitResult, FileVisitor, Files, Path, Paths, SimpleFileVisitor}
import java.lang.ProcessBuilder.Redirect
import java.util.Comparator
import scala.annotation.meta.param
import scala.jdk.CollectionConverters._
import scala.sys.SystemProperties

case class ModuleInfo(name: String, ports: Seq[ModuleInfo.Port]) {
  private[svsim] val instanceName = "dut"
}
object ModuleInfo {
  case class Port(name: String, isSettable: Boolean = false, isGettable: Boolean = false) {
    assert(name.matches("^[a-zA-Z0-9\\-_]*$"))
  }
}

object Workspace {
  val testbenchModuleName: String = "svsimTestbench"

  private class FileDeleter extends SimpleFileVisitor[Path] {
    override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = {
      Files.delete(file)
      FileVisitResult.CONTINUE
    }
    override def postVisitDirectory(dir: Path, ioe: java.io.IOException): FileVisitResult = {
      try {
        Files.delete(dir)
      } catch {
        // This is best effort, leave directories if we can't delete them due to things like FUSE.
        case _: java.nio.file.DirectoryNotEmptyException => ()
      }
      FileVisitResult.CONTINUE
    }
  }
}
final class Workspace(
  path: String,
  /** The prefix for the working directory used when invoking `compile`
    */
  val workingDirectoryPrefix: String = "workdir"
) {

  val absolutePath =
    if (Paths.get(path).isAbsolute())
      path
    else
      s"${System.getProperty("user.dir")}/$path"

  /** A directory where the user can store additional artifacts which are relevant to the primary sources (for instance, artifacts related to the generation of primary sources). These artifacts have no impact on the simulation, but it may be useful to group them with the other files generated by svsim for debugging purposes.
    */
  val supportArtifactsPath = s"$absolutePath/support-artifacts"

  /** The directory containing user-provided source files used to compile the simulation when `compile` is called.
    */
  val primarySourcesPath = s"$absolutePath/primary-sources"

  /** The directory containing code generated when calling `generateAdditionalSources`
    */
  val generatedSourcesPath = s"$absolutePath/generated-sources"

  private var _moduleInfo: Option[ModuleInfo] = None

  def reset() = {
    _moduleInfo = None

    // Create a path type object from the absolute path string.
    val absolutePathObject = Paths.get(absolutePath)
    if (Files.exists(absolutePathObject)) {
      Files.walkFileTree(absolutePathObject, new Workspace.FileDeleter)
    }

    val pathsToCreate = Seq(
      supportArtifactsPath,
      primarySourcesPath,
      generatedSourcesPath
    )
    pathsToCreate.map(new File(_)).foreach(_.mkdirs())
  }

  private def copyResource(klass: Class[_], name: String, targetDirectory: String) = {
    val inputStream = klass.getResourceAsStream(name)
    if (inputStream == null) throw new java.io.FileNotFoundException(name)
    val file = new java.io.File(targetDirectory, name.split("/").last)
    val outputStream = new java.io.FileOutputStream(file)
    Iterator
      .continually(inputStream.read)
      .takeWhile(_ != -1)
      .foreach(outputStream.write)
    inputStream.close()
    outputStream.close()
  }

  /** A helper method which copies the specified resource into the primary sources directory.
    */
  def addPrimarySourceFromResource(klass: Class[_], name: String) = {
    copyResource(klass, name, primarySourcesPath)
  }

  /** `svsim` elaboration simply stores the provided `ModuleInfo` for use by the `compile` method. The idea is that packages that actually do elaboration (like Chisel) will add an overload of this method in an implicit class that then calls this method with the appropriate `ModuleInfo`.
    */
  def elaborate(moduleInfo: ModuleInfo) = {
    assert(_moduleInfo.isEmpty)
    _moduleInfo = Some(moduleInfo)
  }

  /** Generate additional sources necessary for simulating the module.
    */
  //format: off
  final def generateAdditionalSources() = {
    val dut = _moduleInfo.get
    val ports = dut.ports.zipWithIndex

    val systemVerilogTestbenchWriter = new LineWriter(s"$generatedSourcesPath/testbench.sv")
    try {
      val l = systemVerilogTestbenchWriter

      l("module ", Workspace.testbenchModuleName, ";")
      for ((port, index) <- ports) {
        if (port.isSettable) {
      l("  reg  [$bits(", dut.instanceName, ".",  port.name, ")-1:0] ", port.name, " = '0;")
        } else {
      l("  wire [$bits(", dut.instanceName, ".",  port.name, ")-1:0] ", port.name, ";")
        }
      }
      l()
      l(dut.name, " ", dut.instanceName, " (")
      for ((port, index) <- ports) {
      l("    .", port.name, "(", port.name, ")", if (index != ports.length - 1) "," else "")
      }
      l(");")
      l()
      l("  import \"DPI-C\" context function void initTestBenchScope();")
      l("  initial")
      l("    initTestBenchScope();")
      for ((port, index) <- ports) {
      l("  // Port ", index.toHexString, ": ", port.name)
      l("  export \"DPI-C\" function getBitWidthImpl_", port.name, ";")
      l("  function void getBitWidthImpl_", port.name, ";")
      l("    output int value;")
      l("    value", " = $bits(", dut.instanceName, ".",  port.name, ");")
      l("  endfunction")
        if (port.isSettable) {
      l("  export \"DPI-C\" function setBitsImpl_", port.name, ";")
      l("  function void setBitsImpl_", port.name, ";")
      l("    input bit [$bits(", dut.instanceName, ".",  port.name, ")-1:0] value_", port.name, ";")
      l("    ", port.name, " = value_", port.name, ";")
      l("  endfunction")
        }
        if (port.isGettable) {
      l("  export \"DPI-C\" function getBitsImpl_", port.name, ";")
      l("  function void getBitsImpl_", port.name, ";")
      l("    output bit [$bits(", dut.instanceName, ".",  port.name, ")-1:0] value_", port.name, ";")
      l("    value_", port.name, " = ", port.name, ";")
      l("  endfunction")
        }
      l()
      }
      l("  // Simulation")
      l("  import \"DPI-C\" context task simulation_body();")
      l("  enum {INIT, RUN, DONE} simulationState = INIT;")
      l("  initial")
      l("    simulationState = RUN;")
      l("  always @(simulationState) begin")
      l("    if (simulationState == RUN) begin")
      l("      simulation_body();")
      l("      simulationState = DONE;")
      l("    end")
      l("  end")
      l("  import \"DPI-C\" context task simulation_final();")
      l("  final")
      l("    simulation_final();")
      l("  `ifdef ", Backend.HarnessCompilationFlags.supportsDelayInPublicFunctions)
      l("  export \"DPI-C\" task run_simulation;")
      l("  task run_simulation;")
      l("    input int timesteps;")
      l("    output int finish;")
      l("    #timesteps;")
      l("    finish = 0;")
      l("  endtask")
      l("  `else")
      l("  import \"DPI-C\" function void run_simulation(input int timesteps, output int done);")
      l("  `endif")
      l()

      l("  // Tracing")
      l("  int traceSupported = 0;")
      l("  export \"DPI-C\" function simulation_initializeTrace;")
      l("  function void simulation_initializeTrace;")
      l("    input string traceFilePath;")
      l("    `ifdef SVSIM_ENABLE_VCD_TRACING_SUPPORT")
      l("      $dumpfile({traceFilePath,\".vcd\"});")
      l("      $dumpvars(0, ", dut.instanceName,");")
      l("      traceSupported = 1;")
      l("    `endif")
      l("    `ifdef SVSIM_ENABLE_VPD_TRACING_SUPPORT")
      l("      $vcdplusfile({traceFilePath,\".vpd\"});")
      l("      $dumpvars(0, ", dut.instanceName,");")
      l("      $vcdpluson(0, ", dut.instanceName,");")
      l("      traceSupported = 1;")
      l("    `endif")
      l("    `ifdef SVSIM_ENABLE_FSDB_TRACING_SUPPORT")
      l("      $fsdbDumpfile({traceFilePath,\".fsdb\"});")
      l("      $fsdbDumpvars(0, ", dut.instanceName,");")
      l("      traceSupported = 1;")
      l("    `endif")
      l("  endfunction")
      l("  export \"DPI-C\" function simulation_enableTrace;")
      l("  function void simulation_enableTrace;")
      l("    output int success;")
      l("    success = traceSupported;")
      l("    `ifdef SVSIM_ENABLE_VCD_TRACING_SUPPORT")
      l("    $dumpon;")
      l("    `elsif SVSIM_ENABLE_VPD_TRACING_SUPPORT")
      l("    $dumpon;")
      l("    `endif")
      l("    `ifdef SVSIM_ENABLE_FSDB_TRACING_SUPPORT")
      l("    $fsdbDumpon;")
      l("    `endif")
      l("  endfunction")
      l("  export \"DPI-C\" function simulation_disableTrace;")
      l("  function void simulation_disableTrace;")
      l("    output int success;")
      l("    success = traceSupported;")
      l("    `ifdef SVSIM_ENABLE_VCD_TRACING_SUPPORT")
      l("    $dumpoff;")
      l("    `elsif SVSIM_ENABLE_VPD_TRACING_SUPPORT")
      l("    $dumpoff;")
      l("    `endif")
      l("    `ifdef SVSIM_ENABLE_FSDB_TRACING_SUPPORT")
      l("    $fsdbDumpoff;")
      l("    `endif")
      l("  endfunction")
      l()
      l("endmodule")
    } finally {
      // `BufferedWriter` closes the underlying `FileWriter` when closed.
      systemVerilogTestbenchWriter.close()
    }

    // This object creates a wrapper function for exported DPI functions to
    // properly set the scope to testbench top before calling DPI functions.
    object CreateFunctionForPort {
      def createGetBits(portName: String): String = {
        f"""void getBits_$portName(svBitVecVal* result) {
           svScope prev = setScopeToTestBench();
           getBitsImpl_$portName(result);
           svSetScope(prev);
        }"""
      }

      def createGetBitWidth(portName: String): String = {
        f"""void getBitWidth_$portName(int* result) {
           svScope prev = setScopeToTestBench();
           getBitWidthImpl_$portName(result);
           svSetScope(prev);
        }"""
      }
      def createSetBits(portName: String): String = {
        f"""void setBits_$portName(const svBitVecVal* data) {
           svScope prev = setScopeToTestBench();
           setBitsImpl_$portName(data);
           svSetScope(prev);
        }"""
      }
    }

    val cDPIBridgeWriter = new LineWriter(s"$generatedSourcesPath/c-dpi-bridge.cpp")
    try {
      val l = cDPIBridgeWriter
      l("#include <stdint.h>")
      l()
      l("#ifdef SVSIM_ENABLE_VERILATOR_SUPPORT")
      l("#include \"verilated-sources/VsvsimTestbench__Dpi.h\"")
      l("#endif")
      l("#ifdef SVSIM_ENABLE_VCS_SUPPORT")
      l("#include \"vc_hdrs.h\"")
      l("#endif")
      l()
      l("extern \"C\" {")
      l(" svScope setScopeToTestBench();")
      for ((port, index) <- ports) {
      l(CreateFunctionForPort.createGetBitWidth(port.name))
        if (port.isGettable) {
      l(CreateFunctionForPort.createGetBits(port.name))
        }
        if (port.isSettable) {
      l(CreateFunctionForPort.createSetBits(port.name))
        }
      }
      l()
      l("int port_getter(int id, int *bitWidth, void (**getter)(uint8_t*)) {")
      l("  switch (id) {")
      for ((port, index) <- ports.filter(_._1.isGettable)) {
      l("    case ", index.toString(), ": // ", port.name)
      l("      getBitWidth_", port.name, "(bitWidth);")
      l("      *getter = (void(*)(uint8_t*))getBits_", port.name, ";")
      l("      return 0;")
      }
      l("    default:")
      l("      return -1;")
      l("  }")
      l("}")
      l()
      l("int port_setter(int id, int *bitWidth, void (**setter)(const uint8_t*)) {")
      l("  switch (id) {")
      for ((port, index) <- ports.filter(_._1.isSettable)) {
      l("    case ", index.toString(), ": // ", port.name)
      l("      getBitWidth_", port.name, "(bitWidth);")
      l("      *setter = (void(*)(const uint8_t*))setBits_", port.name, ";")
      l("      return 0;")
      }
      l("    default:")
      l("      return -1;")
      l("  }")
      l("}")
      l()
      l("} // extern \"C\"")
      l()

    } finally {
      cDPIBridgeWriter.close()
    }

    copyResource(this.getClass, "/simulation-driver.cpp", generatedSourcesPath)
  }
  //format: on

  /** Compiles the simulation using the specified backend.
    *
    * @param outputTag A string which will be used to tag the output directory. This enables compiling and simulating the same workspace with multiple backends.
    */
  def compile[T <: Backend](
    backend: T
  )(
    workingDirectoryTag:              String,
    commonSettings:                   CommonCompilationSettings,
    backendSpecificSettings:          backend.CompilationSettings,
    customSimulationWorkingDirectory: Option[String],
    verbose:                          Boolean
  ): Simulation = {
    val moduleInfo = _moduleInfo.get
    val workingDirectoryPath = s"$absolutePath/$workingDirectoryPrefix-$workingDirectoryTag"
    val workingDirectory = new File(workingDirectoryPath)
    workingDirectory.mkdir()

    val parameters = backend.generateParameters(
      outputBinaryName = "simulation",
      topModuleName = Workspace.testbenchModuleName,
      additionalHeaderPaths = Seq(workingDirectoryPath),
      commonSettings = commonSettings,
      backendSpecificSettings = backendSpecificSettings
    )

    // Find all the source files by walking the build directory.  The behavior
    // of which directories are visited and which files are included is
    // controlled by fields in `CommonCompilationSettings`.
    val sourceFiles = scala.collection.mutable.ArrayBuffer.empty[String]
    val fileFilter: PartialFunction[File, Boolean] = commonSettings.fileFilter.orElse { case _ => true }
    val directoryFilter: PartialFunction[File, Boolean] = commonSettings.directoryFilter.orElse { case _: File =>
      true
    }
    class DirectoryVisitor extends FileVisitor[Path] {

      override def visitFile(file: Path, attrs: BasicFileAttributes): FileVisitResult = {
        if (fileFilter(file.toFile))
          sourceFiles += workingDirectory.toPath().relativize(file).toString()
        FileVisitResult.CONTINUE
      }

      override def preVisitDirectory(dir: Path, attrs: BasicFileAttributes): FileVisitResult = {
        if (directoryFilter(dir.toFile))
          FileVisitResult.CONTINUE
        else
          FileVisitResult.SKIP_SUBTREE
      }

      override def postVisitDirectory(dir: Path, ioe: java.io.IOException): FileVisitResult = {
        FileVisitResult.CONTINUE
      }

      override def visitFileFailed(file: Path, ioe: java.io.IOException): FileVisitResult = {
        throw ioe
      }

    }

    for (dir <- Seq(primarySourcesPath, generatedSourcesPath)) {
      Files.walkFileTree(Paths.get(dir), new DirectoryVisitor)
    }

    val traceFileStem = (backendSpecificSettings match {
      case s: verilator.Backend.CompilationSettings =>
        s.traceStyle.collectFirst {
          case verilator.Backend.CompilationSettings.TraceStyle.Vcd(_, filename: String) if filename.nonEmpty =>
            filename.stripSuffix(".vcd")
        }
      case _ => None
    }).getOrElse(s"$workingDirectoryPath/trace")
    val simulationEnvironment = Seq(
      "SVSIM_SIMULATION_LOG" -> s"$workingDirectoryPath/simulation-log.txt",
      // The simulation driver appends the appropriate extension to the file path
      "SVSIM_SIMULATION_TRACE" -> traceFileStem
    ) ++ parameters.simulationInvocation.environment

    // Emit Makefile for debugging (will be emitted even if compile fails)
    val makefileWriter = new LineWriter(s"$workingDirectoryPath/Makefile")
    try {
      val l = makefileWriter
      l("# This Makefile enables lightweight debugging of `svsim` tests. ")
      l("# To rebuild the simulation run `make simulation` in this directory.")
      l("# To replay the simulation run `make replay` in this directory.")
      l(
        "# Changes to `generated-sources` and `primary-sources` will be picked up when running `make replay` and `make simulation`. This is useful for debugging issues. You can also freely add, remove or change any of the arguments to the backend or simulation in the targets below."
      )
      l()
      l(".PHONY: clean simulation replay")
      l()
      //format: off
      // For this debug flow, we rebuild the simulation from scratch every time, to avoid issues if the simulation was originally compiled in a different environment, like using SiFive's `wake`.
      l("clean:")
      // Add check if OS is windows, since the command syntax is different
      if (System.getProperty("os.name").toLowerCase.contains("win")) {
        l("\tfor /f \"delims=\" %i in ('dir /b /a-d ^| findstr /v Makefile ^| findstr /v execution-script.txt') do del \"%i\"")
        l("\tfor /d %i in (*) do rmdir /s /q \"%i\"")
      } else {
        l("\tls . | grep -v Makefile | grep -v execution-script.txt | xargs rm -rf")
      }
      l()
      l("simulation: clean")
      l("\t$(compilerEnvironment) \\")
      l("\t", parameters.compilerPath, " \\")
      for (argument <- parameters.compilerInvocation.arguments) {
        val sanitizedArugment = argument
          .replace("$", "$$")
          .replace("'", "'\\''")
          .replace(workingDirectoryPath, "$(shell pwd)")
      l("\t\t'", sanitizedArugment, "' \\")
      }
      l("\t\t$(sourcefiles)")
      l()
      l("replay: simulation")
      val executionScriptPath = "$(shell pwd)/execution-script.txt"
      customSimulationWorkingDirectory match {
        case None =>
        case Some(value) => {
          // Calculate relative path, to avoid being broken by wake's directory shenanigans
          val relativePath = workingDirectory.toPath().relativize(Paths.get(value)).toString()
          l("\tcd ", relativePath, " && \\")
        }
      }
      l("\tcat ", executionScriptPath, " | { grep '^#' || true; } && \\")
      l("\tcat ", executionScriptPath, " | sed -n 's/^[0-9]*> \\(.*\\)/\\1/p' | \\")
      l("\t\t$(simulationEnvironment) $(shell pwd)/simulation \\")
      for (argument <- parameters.simulationInvocation.arguments) {
      l("\t\t\t'", argument.replace("$", "$$"), "' \\")
      }
      l()
      l("sourcefiles = \\")
      for ((sourceFile, index) <- sourceFiles.zipWithIndex) {
      l("\t'", sourceFile, "'", if (index != sourceFiles.length - 1) " \\" else "")
      }
      l()
      l("compilerEnvironment = \\")
      for (((name, value), index) <- parameters.compilerInvocation.environment.zipWithIndex) {
      l("\t", name, "=", value, if (index != parameters.compilerInvocation.environment.length - 1) " \\" else "")
      }
      l()
      l("simulationEnvironment = \\")
      for (((name, value), index) <- simulationEnvironment.zipWithIndex) {
        val sanitizedValue = value.replace(workingDirectoryPath, "$(shell pwd)")
      l("\t", name, "=", sanitizedValue, if (index != simulationEnvironment.length - 1) " \\" else "")
      }
      l()
      //format: on
    } finally {
      makefileWriter.close()
    }

    /**
      * Use the generated Makefile to compile the simulation, since this exercises the Makefile codepath and makes it less likely that we will break `make replay`.
      */
    val processBuilder = new ProcessBuilder("make", "-C", workingDirectoryPath, "simulation")
    processBuilder.redirectErrorStream(true)
    val process = processBuilder.start()
    @scala.annotation.nowarn(
      "msg=Use `scala.jdk.CollectionConverters` instead"
    )
    def readLogLines() = {
      val sourceLocationRegex = "[\\./]*generated-sources/".r
      new BufferedReader(new InputStreamReader(process.getInputStream()))
        .lines()
        .map(sourceLocationRegex.replaceFirstIn(_, ""))
        .map { line =>
          if (verbose) {
            println(line)
          }
          line
        }
        .iterator()
        .asScala
        .toSeq
    }
    val compilationLogLines = readLogLines()
    process.waitFor()
    val compilationLogWriter = new PrintWriter(
      new BufferedWriter(
        new FileWriter(new File(s"$workingDirectoryPath/compilation-log.txt"))
      )
    )
    compilationLogLines.foreach(compilationLogWriter.println)
    compilationLogWriter.close()
    if (process.exitValue() != 0) {
      throw new Exception(compilationLogLines.mkString("\n"))
    }

    new Simulation(
      executableName = "simulation",
      settings = Simulation.Settings(
        customWorkingDirectory = customSimulationWorkingDirectory,
        arguments = parameters.simulationInvocation.arguments,
        environment = simulationEnvironment.toMap
      ),
      workingDirectoryPath = workingDirectoryPath,
      moduleInfo = moduleInfo
    )
  }

}

/** A micro-DSL for writing files.
  */
private class LineWriter(path: String) {
  private val wrapped = new BufferedWriter(new FileWriter(path, false))
  def apply(components: String*) = {
    components.foreach(wrapped.write)
    wrapped.newLine()
  }
  def close() = wrapped.close()
}
